# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
import os
import json
from torch.utils.data import Dataset,DataLoader
from torchvision import datasets, transforms
from torchvision.datasets.folder import ImageFolder, default_loader

from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.data import create_transform

from PIL import Image

class INatDataset(ImageFolder):
    def __init__(self, root, train=True, year=2018, transform=None, target_transform=None,
                 category='name', loader=default_loader):
        self.transform = transform
        self.loader = loader
        self.target_transform = target_transform
        self.year = year
        # assert category in ['kingdom','phylum','class','order','supercategory','family','genus','name']
        path_json = os.path.join(root, f'{"train" if train else "val"}{year}.json')
        with open(path_json) as json_file:
            data = json.load(json_file)

        with open(os.path.join(root, 'categories.json')) as json_file:
            data_catg = json.load(json_file)

        path_json_for_targeter = os.path.join(root, f"train{year}.json")

        with open(path_json_for_targeter) as json_file:
            data_for_targeter = json.load(json_file)

        targeter = {}
        indexer = 0
        for elem in data_for_targeter['annotations']:
            king = []
            king.append(data_catg[int(elem['category_id'])][category])
            if king[0] not in targeter.keys():
                targeter[king[0]] = indexer
                indexer += 1
        self.nb_classes = len(targeter)

        self.samples = []
        for elem in data['images']:
            cut = elem['file_name'].split('/')
            target_current = int(cut[2])
            path_current = os.path.join(root, cut[0], cut[2], cut[3])

            categors = data_catg[target_current]
            target_current_true = targeter[categors[category]]
            self.samples.append((path_current, target_current_true))

    # __getitem__ and __len__ inherited from ImageFolder

class ImageDataset(Dataset):
    ### 读取训练json
    def __init__(self, data_json,is_train,args):
        super(ImageDataset, self).__init__()
        self.trans = build_transform(False, args)
        with open(data_json) as f:
            self.data_json = json.load(f)
        if is_train:
            self.data_json = self.data_json['train']
        else:
            self.data_json = self.data_json['valid']

    def __len__(self):
        return len(self.data_json)

    def __getitem__(self, index):
        pattern_name,model_name,label = self.data_json[index]  # name: dataset/train/oript_searchpt.jpg
        pattern_img = self.trans(Image.open(pattern_name).convert('RGB'))
        model_img = self.trans(Image.open(model_name).convert('RGB'))
        return model_img, pattern_img,label

import numpy as np

class TestPatternDataset(Dataset):
    ### 测试花型集
    def __init__(self, pattern_dir,args):
        super(TestPatternDataset, self).__init__()
        self.trans = build_transform(False, args)
        self.pattern_datadir = pattern_dir
        self.patternlist = os.listdir(pattern_dir)

    def __len__(self):
        return len(self.patternlist)

    def __getitem__(self, index):
        pattern_name = os.path.join(self.pattern_datadir, self.patternlist[index]) # name: dataset/train/oript_searchpt.jpg
        pattern_img = self.trans(Image.open(pattern_name).convert('RGB'))
        # transforms.ToPILImage()(pattern_img).convert("RGB").save('./pattern_img.png')
        return pattern_name,pattern_img

# class ImageDataset(Dataset):
#     ### 读取花型文件夹作为名字
#     def __init__(self, model_datadir,pattern_datadir,is_train,args):
#         super(ImageDataset, self).__init__()
#         self.model_datadir = model_datadir
#         self.pattern_datadir = pattern_datadir
#         self.patternlist = os.listdir(pattern_datadir)
#         self.trans = build_transform(is_train, args)
#
#     def __len__(self):
#         return len(self.patternlist)
#
#     def __getitem__(self, index):
#         pattern_name = self.patternlist[index]  # name: oript_searchpt
#         name,_IMAGE_EXTENSION = pattern_name.split('.')
#         oript, searchpt = name.split('_')
#         pattern_img = self.trans(Image.open(os.path.join(self.pattern_datadir,pattern_name)).convert('RGB'))
#         model_img = self.trans(Image.open(os.path.join(self.model_datadir, oript + '.png')).convert('RGB'))
#         if oript == searchpt:
#             label = 0  #属于同一个类，Y=0
#         else:
#             label = 1  #不属于同一个类，Y=1
#         return model_img, pattern_img,label


def build_transform(is_train, args):
    resize_im = args.input_size > 32
    if is_train:
        # this should always dispatch to transforms_imagenet_train
        transform = create_transform(
            input_size=args.input_size,
            is_training=True,
            color_jitter=args.color_jitter,
            auto_augment=args.aa,
            interpolation=args.train_interpolation,
            re_prob=args.reprob,
            re_mode=args.remode,
            re_count=args.recount,
        )
        if not resize_im:
            # replace RandomResizedCropAndInterpolation with
            # RandomCrop
            transform.transforms[0] = transforms.RandomCrop(
                args.input_size, padding=4)
        return transform

    t = []
    if resize_im:
        size = int((256 / 224) * args.input_size)
        t.append(
            transforms.Resize(size, interpolation=3),  # to maintain same ratio w.r.t. 224 images
        )
        t.append(transforms.CenterCrop(args.input_size))

    t.append(transforms.ToTensor())
    # t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD))
    return transforms.Compose(t)
